{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neo4j import GraphDatabase\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "uri = \"\"\n",
    "username = \"\"\n",
    "password = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>id</th>\n",
       "      <th>url</th>\n",
       "      <th>rate</th>\n",
       "      <th>length</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>我是路人甲</td>\n",
       "      <td>0</td>\n",
       "      <td>http://movie.douban.com/subject/25746375/</td>\n",
       "      <td>7.4</td>\n",
       "      <td>134.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>铁拳</td>\n",
       "      <td>1</td>\n",
       "      <td>http://movie.douban.com/subject/5446197/</td>\n",
       "      <td>7.1</td>\n",
       "      <td>123.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>我们梦中见</td>\n",
       "      <td>2</td>\n",
       "      <td>http://movie.douban.com/subject/25885212/</td>\n",
       "      <td>7.6</td>\n",
       "      <td>92.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>少年透明人</td>\n",
       "      <td>3</td>\n",
       "      <td>http://movie.douban.com/subject/25728581/</td>\n",
       "      <td>6.6</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>撒迦利亚</td>\n",
       "      <td>4</td>\n",
       "      <td>http://movie.douban.com/subject/5156558/</td>\n",
       "      <td>6.0</td>\n",
       "      <td>95.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4582</th>\n",
       "      <td>幽灵世界</td>\n",
       "      <td>4582</td>\n",
       "      <td>http://movie.douban.com/subject/1304868/</td>\n",
       "      <td>7.9</td>\n",
       "      <td>111.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4583</th>\n",
       "      <td>吮拇指的人</td>\n",
       "      <td>4583</td>\n",
       "      <td>http://movie.douban.com/subject/1422954/</td>\n",
       "      <td>7.2</td>\n",
       "      <td>96.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4584</th>\n",
       "      <td>歌舞青春3：毕业季</td>\n",
       "      <td>4584</td>\n",
       "      <td>http://movie.douban.com/subject/2215609/</td>\n",
       "      <td>7.3</td>\n",
       "      <td>112.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4585</th>\n",
       "      <td>大河恋</td>\n",
       "      <td>4585</td>\n",
       "      <td>http://movie.douban.com/subject/1292718/</td>\n",
       "      <td>8.4</td>\n",
       "      <td>123.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4586</th>\n",
       "      <td>青春年少</td>\n",
       "      <td>4586</td>\n",
       "      <td>http://movie.douban.com/subject/1296835/</td>\n",
       "      <td>7.9</td>\n",
       "      <td>93.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4587 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           name    id                                        url  rate  length\n",
       "0         我是路人甲     0  http://movie.douban.com/subject/25746375/   7.4   134.0\n",
       "1            铁拳     1   http://movie.douban.com/subject/5446197/   7.1   123.0\n",
       "2         我们梦中见     2  http://movie.douban.com/subject/25885212/   7.6    92.0\n",
       "3         少年透明人     3  http://movie.douban.com/subject/25728581/   6.6   100.0\n",
       "4          撒迦利亚     4   http://movie.douban.com/subject/5156558/   6.0    95.0\n",
       "...         ...   ...                                        ...   ...     ...\n",
       "4582       幽灵世界  4582   http://movie.douban.com/subject/1304868/   7.9   111.0\n",
       "4583      吮拇指的人  4583   http://movie.douban.com/subject/1422954/   7.2    96.0\n",
       "4584  歌舞青春3：毕业季  4584   http://movie.douban.com/subject/2215609/   7.3   112.0\n",
       "4585        大河恋  4585   http://movie.douban.com/subject/1292718/   8.4   123.0\n",
       "4586       青春年少  4586   http://movie.douban.com/subject/1296835/   7.9    93.0\n",
       "\n",
       "[4587 rows x 5 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv(\"./data/entities/movies.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver = GraphDatabase.driver(uri, auth=(username, password))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_query(tx, query, parameters=None):\n",
    "    tx.run(query, parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def special_word(strs):\n",
    "    rgx = [\"'\",'=~','~', '*', '%', '^', '<=', '>=', '=', '<>', '!=', 'XOR', 'OR', '|', '\\\\']\n",
    "    for rg in rgx:\n",
    "        strs = strs.replace(rg, ' ')\n",
    "    strs.replace('\\'', '')\n",
    "    return strs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1', '2', '3']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "['1','2','3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1,2,3'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "','.join(['1','2','3'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with driver.session() as session:\n",
    "    for row in tqdm(range(len(df))):\n",
    "        sx_dic = {}\n",
    "        for sx in ['id','url','rate','length']:\n",
    "            if df.loc[row][sx]:\n",
    "                if sx == 'id':\n",
    "                    sx_dic['code'] = df.loc[row][sx]\n",
    "                else:\n",
    "                    sx_dic[sx] = df.loc[row][sx]\n",
    "        query = \"CREATE (p:\"\n",
    "        query += \"Movie {\"\n",
    "        query += \"name:'\" + special_word(str(df.loc[row]['name'])) + \"',\"\n",
    "        query+= ','.join([f\"{k} :  '{v}'\" for k, v in sx_dic.items()])\n",
    "        query += \"})\"\n",
    "        try:\n",
    "            # 执行查询\n",
    "            session.write_transaction(execute_query, query)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "        \n",
    "        # CREATE(<node_name>:<label_name> {<property1_name>:<property1_value>,...,<propertyn_name>:<propertyn_value>) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_node(name_str):\n",
    "    df = pd.read_csv('./data/entities/'+name_str+'.csv')\n",
    "    with driver.session() as session:\n",
    "        for row in tqdm(range(len(df))):\n",
    "            query = \"CREATE (p:\" + name_str.rstrip('s').title()\n",
    "            query += \" {name:'\" + special_word(str(df.loc[row]['name'])) + \"',\"\n",
    "            query += \"code:'\" + str(df.loc[row]['id'])\n",
    "            query += \"'})\"\n",
    "            try:\n",
    "                # 执行查询\n",
    "                session.write_transaction(execute_query, query)\n",
    "            except Exception as e:\n",
    "                print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_node('years')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_node('districts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_node('mv_classes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_node('languages')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_node('people')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "relation_df = pd.read_csv('./data/relations/triple.csv')\n",
    "relation_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for row in tqdm(range(len(relation_df))):\n",
    "    start_node = str(relation_df.loc[row]['Entity1'])\n",
    "    end_node = str(relation_df.loc[row]['Entity2'])           \n",
    "    # 查找已存在的节点\n",
    "    with driver.session() as session:\n",
    "        s1 = \" (a:Movie {code: $code1})\"\n",
    "        if relation_df.loc[row]['Relation'] in ['导演','编剧','演员']:\n",
    "            s2 = \" (b:People {code: $code2})\"\n",
    "        elif relation_df.loc[row]['Relation'] == '上映时间':\n",
    "            s2 = \" (b:Year {code: $code2})\"\n",
    "        elif relation_df.loc[row]['Relation'] == '上映地区':\n",
    "            s2 = \" (b:District {code: $code2})\"\n",
    "        elif relation_df.loc[row]['Relation'] == '类型':\n",
    "            s2 = \" (b:Mv_Classe {code: $code2})\"\n",
    "        elif relation_df.loc[row]['Relation'] == '语言':\n",
    "            s2 = \" (b:Language {code: $code2})\"\n",
    "\n",
    "    # 创建关系\n",
    "    with driver.session() as session:\n",
    "        try:\n",
    "            session.run(\"MATCH\"+s1+\",\"+s2+\"\"\n",
    "                        \"CREATE (a) - [:\"+str(relation_df.loc[row]['Relation'])+\"] -> (b)\",\n",
    "                        code1=start_node,code2=end_node)\n",
    "        except Exception as e:\n",
    "            print(\"create relation error:\" , e)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
